import { useState, useEffect, useCallback } from 'react';
import { ollamaService, OllamaModel } from '../services/OllamaService';
import { CommandResult } from '../services/TerminalService';

interface ModelOptions {
  format?: 'json' | 'text';
  keepAlive?: string;
  insecure?: boolean;
  noWordWrap?: boolean;
}

interface ServerOptions {
  debug?: boolean;
  maxModelsPerGpu?: number;
  maxQueueRequests?: number;
  parallelRequests?: number;
  scheduledSpread?: boolean;
  flashAttention?: boolean;
}

export const useOllama = () => {
  const [models, setModels] = useState<OllamaModel[]>([]);
  const [isServerRunning, setIsServerRunning] = useState(false);
  const [isLoading, setIsLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);

  useEffect(() => {
    const subscription = ollamaService.serverStatus$.subscribe(
      status => setIsServerRunning(status)
    );
    return () => subscription.unsubscribe();
  }, []);

  useEffect(() => {
    if (isServerRunning) {
      refreshModels();
    } else {
      setModels([]);
    }
  }, [isServerRunning]);

  const refreshModels = useCallback(async () => {
    setIsLoading(true);
    setError(null);
    try {
      const modelList = await ollamaService.listModels();
      setModels(modelList);
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Failed to fetch models');
    } finally {
      setIsLoading(false);
    }
  }, []);

  const startServer = async (options?: ServerOptions): Promise<CommandResult> => {
    setIsLoading(true);
    setError(null);
    try {
      return await ollamaService.startServer(options);
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Failed to start server');
      throw err;
    } finally {
      setIsLoading(false);
    }
  };

  const stopServer = async (): Promise<CommandResult> => {
    setIsLoading(true);
    setError(null);
    try {
      return await ollamaService.stopServer();
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Failed to stop server');
      throw err;
    } finally {
      setIsLoading(false);
    }
  };

  const createModel = async (name: string, modelfile: string): Promise<CommandResult> => {
    setIsLoading(true);
    setError(null);
    try {
      const result = await ollamaService.createModel(name, modelfile);
      await refreshModels();
      return result;
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Failed to create model');
      throw err;
    } finally {
      setIsLoading(false);
    }
  };

  const deleteModel = async (name: string): Promise<CommandResult> => {
    setIsLoading(true);
    setError(null);
    try {
      const result = await ollamaService.deleteModel(name);
      await refreshModels();
      return result;
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Failed to delete model');
      throw err;
    } finally {
      setIsLoading(false);
    }
  };

  const runModel = async (name: string, prompt: string, options?: ModelOptions): Promise<CommandResult> => {
    setIsLoading(true);
    setError(null);
    try {
      return await ollamaService.runModel(name, prompt, options);
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Failed to run model');
      throw err;
    } finally {
      setIsLoading(false);
    }
  };

  const pullModel = async (name: string, options?: { insecure?: boolean }): Promise<CommandResult> => {
    setIsLoading(true);
    setError(null);
    try {
      const result = await ollamaService.pullModel(name, options);
      await refreshModels();
      return result;
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Failed to pull model');
      throw err;
    } finally {
      setIsLoading(false);
    }
  };

  const pushModel = async (name: string, options?: { insecure?: boolean }): Promise<CommandResult> => {
    setIsLoading(true);
    setError(null);
    try {
      return await ollamaService.pushModel(name, options);
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Failed to push model');
      throw err;
    } finally {
      setIsLoading(false);
    }
  };

  return {
    models,
    isServerRunning,
    isLoading,
    error,
    startServer,
    stopServer,
    createModel,
    deleteModel,
    runModel,
    pullModel,
    pushModel,
    refreshModels,
  };
}; 